# 本脚本用于测试dpa在线线程调度+线程运行的时间消耗
import subprocess
import os
import csv
import re

# 创建结果目录
results_dir = '../results_raw'
os.makedirs(results_dir, exist_ok=True)

# 结果文件
output_file = os.path.join(results_dir, 'dpa_thread_analyze_v2.csv')

# 生成batch_size序列: 2, 16, 128, ..., 4194304
batch_sizes = []
current = 8192
while current <= 8192:
    batch_sizes.append(current)
    current *= 8

# 生成thread_num序列
thread_nums = [1, 2, 4, 8, 16, 32, 64, 128, 253]
# thread_nums = [253]

def convert_to_us(value, unit):
    """将时间单位转换为微秒"""
    unit = unit.lower()
    if unit == 's':
        return value * 1_000_000
    elif unit == 'ms':
        return value * 1_000
    elif unit == 'us':
        return value
    elif unit == 'ns':
        return value / 1000
    else:
        raise ValueError(f"未知单位: {unit}")

# 打开文件以写入结果
with open(output_file, 'w', newline='') as csvfile:
    fieldnames = [
        'op_type',
        'batch_size',
        'thread_num',
        'thread_time',
        'p50_time',
        'p99_time',
        'thread_ratio'
    ]
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    total_runs = len(batch_sizes) * len(thread_nums)
    current_run = 0
    
    for batch_size in batch_sizes:
        for thread_num in thread_nums:
            current_run += 1
            print(f"运行测试 ({current_run}/{total_runs}): batch_size={batch_size}, thread_num={thread_num}")
            
            # 构建命令
            command = [
                './../../bin_dpu/test_operator',
                f'--batch_size={batch_size}',
                f'--thread_num={thread_num}',
                '--device_type=dpa',
                '--op_type=sigrid_hash',
                '--dpa_type=2'
            ]

            # 执行命令并捕获输出
            try:
                result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                        text=True, encoding='utf-8', errors='ignore', timeout=60)
            except subprocess.TimeoutExpired:
                print(f"超时: batch_size={batch_size}, thread_num={thread_num}")
                continue
            except Exception as e:
                print(f"执行错误: {e}")
                continue

            stdout = result.stdout
            stderr = result.stderr
            
            if stderr:
                print(f"错误输出: {stderr}")
            
            # 1. 提取 max_execute_time_per_thread
            thread_time = None
            max_time_pattern = r'max_execute_time_per_thread\s*=\s*([\d.]+)\s*([a-zA-Z]+)'
            for line in stdout.splitlines():
                match = re.search(max_time_pattern, line)
                if match:
                    try:
                        value = float(match.group(1))
                        unit = match.group(2)
                        thread_time = convert_to_us(value, unit)
                        break  # 找到后立即退出循环
                    except (ValueError, TypeError):
                        continue
            
            if thread_time is None:
                print(f"警告: 未找到max_execute_time_per_thread - batch_size={batch_size}, thread_num={thread_num}")
                continue
            
            # 2. 提取时间统计表格
            op_type = None
            p50_time = None
            p99_time = None
            
            # 使用正则表达式匹配表格行
            table_pattern = r'\| (\w+)\s+\| (\d+\.?\d*)\s+(\w+)\s+\| (\d+\.?\d*)\s+(\w+)\s+\|'
            for line in stdout.splitlines():
                match = re.search(table_pattern, line)
                if match:
                    op_type = match.group(1)
                    p50_value = float(match.group(2))
                    p50_unit = match.group(3)
                    p99_value = float(match.group(4))
                    p99_unit = match.group(5)
                    
                    # 转换为微秒
                    p50_time = convert_to_us(p50_value, p50_unit)
                    p99_time = convert_to_us(p99_value, p99_unit)
                    break
            
            if op_type is None or p50_time is None or p99_time is None:
                print(f"警告: 未找到时间统计表格 - batch_size={batch_size}, thread_num={thread_num}")
                continue
            
            # 计算thread_ratio (0-100的整数)
            thread_ratio = int(round((thread_time / p50_time) * 100, 0))
            thread_ratio = max(0, min(100, thread_ratio))
            
            # 写入CSV行
            writer.writerow({
                'op_type': op_type,
                'batch_size': batch_size,
                'thread_num': thread_num,
                'thread_time': thread_time,
                'p50_time': p50_time,
                'p99_time': p99_time,
                'thread_ratio': thread_ratio
            })

print(f"测试完成! 结果已保存到 {output_file}")
